# -*- coding: utf-8 -*-
# @Organization  : 南京畅淼科技有限责任公司
# @Author        : 郭金阳
# @Time          : 2023/4/17 15:17
# @Function      :
import time
import os
import cv2
from log import log
from os.path import join
from datetime import date
import base64
import redis
import yaml
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import json
from subprocess import PIPE, Popen
import re

try:
    with open('./config.yml', 'r', encoding='utf-8') as config_file:
        cfg = yaml.load(config_file, Loader=yaml.FullLoader)
except:
    log.logger.error("配置文件读取失败，请核查")


class VideoSave:
    def __init__(self, init):
        self.camIP = init['cam_iP']
        self.devPort = init['dev_port']
        self.userName = init['username']
        self.passWord = init['password']
        # self.frameRate = init['frame_rate']
        self.videoW = init['video_width']
        self.videoH = init['video_high']
        self.videoFps = init['video_fps']
        self.isSaveVideo = init['is_save_video']
        self.videoSavePath = init['video_save_path']
        self.videoMaxSize = init['video_max_size']
        self.isSaveVideo = init["is_save_video"]
        self.count = 0
        self.curWriting = None  # 当前正在写入本地视频的视频路径
        self.videoWriter = None  # 视频写入类对象
        self.pushFailTime = None  # 用于记录rtmp流推送失败发生时间
        self.writeFailTime = None  # 用于记录视频流写入本地失败发生时间
        self.videoPushUrl = init["video_push_url"]
        self.command = ['./ffmpeg',  # 待执行指令，通过PIPE管道、ffmpeg将处理好的帧推送rtmp流
                        '-y',
                        '-f', 'rawvideo',
                        '-pix_fmt', 'bgr24',
                        '-s', "{}x{}".format(int(self.videoW), int(self.videoH)),
                        '-r', str(self.videoFps),
                        '-i', '-',
                        '-c:v', 'libx264',
                        '-pix_fmt', 'yuv420p',
                        '-preset', 'faster',
                        '-f', 'flv',
                        self.videoPushUrl]

        self.rtmpOrLocal = init["rtmp_or_local"]
        # 连接redis数据库
        self.redis = self.getRedis(init["redis_out_ip"], init["redis_out_port"],
                                   init["redis_out_pwd"])
        self.frame = init["res_task_key_frame"]
        self.box = init["res_task_key_box"]
        self.snap_frame = init["res_task_key_snap_frame"]
        self.pusher = Popen(self.command, stdin=PIPE)

    def convert_base64_to_video(self):
        while True:
            # 从redis中取出检测视频帧
            img_frame_detect = self.redis.hget(self.frame, "1")
            dict_img_detect = json.loads(img_frame_detect)
            img_frame_detect = dict_img_detect['jpg_as_text']
            # 将base64编码的字符串转换为字节流
            decoded_bytes_detect = base64.b64decode(img_frame_detect)
            np_arr_detect = np.frombuffer(decoded_bytes_detect, np.uint8)
            image_detect = cv2.imdecode(np_arr_detect, cv2.COLOR_BGR2RGB)
            frame_detect_box = dict_img_detect['frame_detect_box']
            detect_info_dict = eval(frame_detect_box)
            for keys in detect_info_dict:
                if keys == "-1":
                    for i in detect_info_dict[keys]:
                        det_box = i
                        cv2.rectangle(image_detect, (int(float(det_box[0])), int(float(det_box[1]))),
                                      (int(float(det_box[2])), int(float(det_box[3]))),
                                      (129, 51, 255), 2)
                    continue
                det_box = detect_info_dict[keys]["detBox"]
                ais_ship_name = detect_info_dict[keys]["aisShipName"]
                ais_mmsi = detect_info_dict[keys]["aisMmsi"]
                snap_distance = detect_info_dict[keys]["snapDistance"]
                current_state = detect_info_dict[keys]["currentState"]
                # 已抓拍
                if current_state == 1:
                    cv2.rectangle(image_detect, (int(float(det_box[0])), int(float(det_box[1]))),
                                  (int(float(det_box[2])), int(float(det_box[3]))),
                                  (255, 248, 51), 2)
                # 正在抓拍
                elif current_state == 3:
                    cv2.rectangle(image_detect, (int(float(det_box[0])), int(float(det_box[1]))),
                                  (int(float(det_box[2])), int(float(det_box[3]))),
                                  (51, 255, 54), 2)
                    # 绘制矩形并填充颜色
                    # 设置矩形不透明度
                    alpha = 0.45
                    overlay = image_detect.copy()
                    cv2.rectangle(overlay, (int(float(det_box[0])), int(float(det_box[1]))),
                                  (int(float(det_box[2])), int(float(det_box[3]))), (51, 255, 54), -1)
                    cv2.addWeighted(overlay, alpha, image_detect, 1 - alpha, 0, image_detect)
                    cv2.line(image_detect, (640, 719), ((int(float(det_box[0])) +
                                                         int(float(det_box[2]))) // 2, int(float(det_box[3]))),
                             (51, 255, 54), 5)
                # 未抓拍
                elif current_state == 0:
                    cv2.rectangle(image_detect, (int(float(det_box[0])), int(float(det_box[1]))),
                                  (int(float(det_box[2])), int(float(det_box[3]))),
                                  (51, 255, 255), 2)
                # 图片转换（cv2 -> pil）
                cv2_img = cv2.cvtColor(image_detect, cv2.COLOR_BGR2RGB)
                pil_img = Image.fromarray(cv2_img)
                # 在图片上添加文字（支持中文）
                draw = ImageDraw.Draw(pil_img)
                font = ImageFont.truetype("simhei.ttf", 15, encoding="utf-8")
                # 已抓拍
                if current_state == 1:
                    draw.text((det_box[0], det_box[1] - 20), f"MMSI-{ais_mmsi}", (51, 248, 255), font=font)
                    draw.text((det_box[0], det_box[1] - 40), f"船名-{ais_ship_name}", (51, 248, 255), font=font)
                    log.logger.debug(f"该船船名为{ais_ship_name}，状态为{current_state}，检测框位置为{det_box}")
                # 正在抓拍
                elif current_state == 3:
                    draw.text((det_box[0], det_box[1] - 20), f"MMSI-{ais_mmsi}", (54, 255, 51), font=font)
                    draw.text((det_box[0], det_box[1] - 40), f"船名-{ais_ship_name}", (54, 255, 51), font=font)
                    draw.text(((det_box[0] + det_box[2]) // 2, (det_box[1] + det_box[3]) // 2), f"抓拍对象",
                              (255, 255, 255), font=font)
                    draw.text(((640 + (int(float(det_box[0])) + int(float(det_box[2]))) // 2) // 2,
                               (719 + int(float(det_box[3]))) // 2), f"抓拍距离-{snap_distance}米",
                              (54, 255, 51), font=font)
                    log.logger.debug(f"该船船名为{ais_ship_name}，状态为{current_state}，检测框位置为{det_box}")
                # 未抓拍
                elif current_state == 0:
                    draw.text((det_box[0], det_box[1] - 20), f"MMSI-{ais_mmsi}", (255, 255, 51), font=font)
                    draw.text((det_box[0], det_box[1] - 40), f"船名-{ais_ship_name}", (255, 255, 51), font=font)
                    log.logger.debug(f"该船船名为{ais_ship_name}，状态为{current_state}，检测框位置为{det_box}")
                # 图片转换（pil -> cv2）
                image_detect = cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
                # 从redis中取出抓拍视频帧
                redis_list_snap = list(map(int, self.redis.hkeys(self.snap_frame)))
                if redis_list_snap:
                    for key_cap in sorted(redis_list_snap, reverse=True):
                        img_frame_capture = self.redis.hget(self.snap_frame,
                                                            key_cap)
                        dict_img = json.loads(img_frame_capture)
                        img_frame_capture = dict_img['frame']
                        statues_capture = dict_img['statues']
                        if statues_capture:
                            log.logger.debug(f"正在抓拍的redis数据{key_cap, statues_capture}")
                            # # 将base64编码的字符串转换为字节流
                            decoded_bytes_capture = base64.b64decode(img_frame_capture)
                            np_arr_capture = np.frombuffer(decoded_bytes_capture, np.uint8)
                            image_capture = cv2.imdecode(np_arr_capture, cv2.IMREAD_COLOR)

                            resized_img = cv2.resize(image_capture, (500, 300), interpolation=cv2.INTER_LINEAR)
                            x_offset = 780
                            y_offset = 420
                            image_detect[y_offset:y_offset + resized_img.shape[0], x_offset:x_offset +
                                                                                            resized_img.shape[
                                                                                                1]] = resized_img
                        # 删除redisB元素
                        else:
                            log.logger.debug(f"删除抓拍过的redis数据{key_cap, statues_capture}")
                            self.redis.hdel(self.snap_frame, key_cap)
                            # 保存视频帧
            if self.rtmpOrLocal:
                self.push_rtmp_video(image_detect)
            else:
                self.push_video(image_detect)
            # self.push_video(image_detect)

    # 推送rtmp流以及将视频保存到本地
    def push_rtmp_video(self, frame):
        log.logger.info('推流线程开始运行')
        # 创建用于执行ffmpeg命令的子进程，用于通过ffmpeg将处理好的帧推送rtmp流
        # try:
        #     self.pusher = Popen(self.command, stdin=PIPE)
        # except:
        #     log.logger.error('推流线程启动失败，请检查ffmpeg程序是否可以成功运行')
        # 如果需要保存视频并且没有视频流写入对象时，创建一个视频流写入对象
        if self.isSaveVideo == 1 and self.videoWriter is None:
            self.videoWriter = self.get_video_writer()
        # 当不需要保存时，将当前视频写入对象释放
        if self.isSaveVideo == 0 and self.videoWriter is not None:
            self.videoWriter.release()
            self.videoWriter = None
            self.curWriting = None
            log.logger.info('停止向本地写入视频')
        # 从通信队列中取出视频流检测线程已经处理好的帧
        frame = cv2.resize(frame, (self.videoW, self.videoH))
        self.push_frame(frame)

    # 将处理好的帧推送rtmp视频流
    def push_frame(self, frame):
        """
        :param frame: 处理好的帧
        """
        try:
            # 将处理好的帧输送到PIPE管道中，并通过ffmpeg推送rtmp流
            self.pusher.stdin.write(frame.tostring())
            self.pushFailTime = None
        except:
            # 当推送rtmp视频流失败时，记录下推送失败发生的时间
            if self.pushFailTime is None:
                self.pushFailTime = time.time()
                log.logger.error('推送rtmp流失败，请检测nginx服务是否启动并检测配置文件中rtmp流的url')
            else:
                # 若推送rtmp流持续失败，每隔1分钟打下log
                curTime = time.time()
                if curTime - self.pushFailTime >= 60:
                    log.logger.error('推送rtmp流失败，请检测nginx服务是否启动并检测配置文件中rtmp流的url')
                    self.pushFailTime = curTime

    # 将视频保存到本地
    def push_video(self, frame):
        # 如果需要保存视频并且没有视频流写入对象时，创建一个视频流写入对象
        if self.isSaveVideo == 1 and self.videoWriter is None:
            self.videoWriter = self.get_video_writer()
        # 当不需要保存时，将当前视频写入对象释放
        if self.isSaveVideo == 0 and self.videoWriter is not None:
            self.videoWriter.release()
            self.videoWriter = None
            self.curWriting = None
            log.logger.info('停止向本地写入视频')
        # frame = cv2.resize(frame, (self.videoW, self.videoH))
        if self.isSaveVideo == 1:
            self.write_frame(frame)

    # 将处理好的帧写入本地
    def write_frame(self, frame):
        """
        :param frame: 处理好的帧
        """
        if self.isSaveVideo == 0:
            return
        # 如果当前正在写入的视频占存大小超过配置文件中所设定的最大占存，就写入到新的视频中
        if os.stat(self.curWriting).st_size / (1000 ** 2) >= self.videoMaxSize:
            self.videoWriter.release()
            self.videoWriter = self.get_video_writer()
        try:
            # 写入视频流到本地
            self.videoWriter.write(frame)
            self.writeFailTime = None
        except:
            # 当写入视频流到本地失败时，记录下写入失败发生的时间
            if self.writeFailTime is None:
                self.writeFailTime = time.time()
                log.logger.error('视频写入本地失败，请检查本地磁盘空间')
            else:
                # 若写入视频流到本地持续失败，每隔1分钟打下log
                curTime = time.time()
                if curTime - self.writeFailTime >= 60:
                    log.logger.error('视频写入本地失败，请检查本地磁盘空间')
                    self.pushFailTime = curTime

    # 创建一个视频写入类的对象
    def get_video_writer(self):
        if not os.path.exists(self.videoSavePath):
            os.mkdir(self.videoSavePath)
        idx = 1
        # 本地写入视频命名格式，年月日加编号
        videoName = '%s__%d.avi' % (date.today(), idx)
        videoPath = join(self.videoSavePath, videoName)
        while os.path.exists(join(self.videoSavePath, videoName)):
            idx += 1
            videoName = '%s__%d.avi' % (date.today(), idx)
            videoPath = join(self.videoSavePath, videoName)
        # 创建一个视频写入类的对象
        # fourcc = cv2.VideoWriter_fourcc(*'MP42')*'mp4v'
        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
        # fourcc = cv2.VideoWriter_fourcc('M', 'J', 'P', 'G')
        writer = cv2.VideoWriter(videoPath, fourcc,
                                 self.videoFps, (self.videoW, self.videoH))
        self.curWriting = videoPath
        log.logger.info('向本地写入视频，路径为%s' % videoPath)
        return writer

    def getRedis(self, host, port, pwd):
        """
        连接redis
        :param host:
        :param port:
        :param pwd:
        :return:
        """
        try:
            pool = redis.ConnectionPool(host=host, port=port, password=pwd, db=0)
            r = redis.Redis(connection_pool=pool)
            log.logger.info("Redis连接成功")
            return r
        except:
            log.logger.info('redis连接失败,正在尝试重连')
            return None


if __name__ == '__main__':
    pass

# pyinstaller -F utils.py
